package com.aiguigu.tree;

import java.util.*;

public class HuffmanCode {
    static StringBuilder stringBuilder=new StringBuilder();
    static Map<Byte,String> huffmanCodes=new HashMap<>();
    public static void main(String[] args) {
        String string = "iskjldhsakjgh sdkjahf";
        byte[] bytes = string.getBytes();
        List<CodeNode> listOfCode = getListOfCode(bytes);
        CodeNode codeNode = HuffmanTree(listOfCode);
        codeNode.preOrder();
       getCode(codeNode);
        System.out.println(huffmanCodes);
    }

    public static List<CodeNode> getListOfCode(byte[] bytes){
        ArrayList<CodeNode> nodes = new ArrayList<>();
        Map<Byte,Integer>counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if (count == null){
                counts.put(b,1);
            }else {
                counts.put(b,count+1);
            }
        }
        for (Map.Entry<Byte,Integer> entry:counts.entrySet()) {
            nodes.add(new CodeNode(entry.getKey(),entry.getValue()));
        }
        return nodes;
    }



    public static CodeNode HuffmanTree(List<CodeNode> codeNodeList ){
        while(codeNodeList.size()>1){
            Collections.sort(codeNodeList);
            CodeNode leftChild = codeNodeList.get(0);
            CodeNode rightChild = codeNodeList.get(1);

            CodeNode parent = new CodeNode(leftChild.weight+rightChild.weight);
            parent.left=leftChild;
            parent.right= rightChild;
            codeNodeList.remove(leftChild);
            codeNodeList.remove(rightChild);
            codeNodeList.add(parent);
        }
        return codeNodeList.get(0);
    }

    public static Map<Byte,String> getCode(CodeNode root){
        if (root == null){
            return null;
        }
       /* getCode(root.left,"0",stringBuilder);
        getCode(root.right,"1",stringBuilder);*/
       getCode(root,"",stringBuilder);
        return huffmanCodes;

    }

    public static void getCode(CodeNode node,String code,StringBuilder stringBuilder){
        StringBuilder stringBuilder2 = new StringBuilder(stringBuilder);
        stringBuilder2.append(code);
        if (node!=null){
            if (node.data ==null){
                getCode(node.left,"0",stringBuilder2);
                getCode(node.right,"1",stringBuilder2);
            }else{
                huffmanCodes.put(node.data,stringBuilder2.toString());
            }
        }
    }

}

class CodeNode implements Comparable<CodeNode>{
    Byte data;
    int weight;
    CodeNode left;
    CodeNode right;

    public CodeNode(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    public CodeNode(int weight) {
        this.weight = weight;
    }

    @Override
    public int compareTo(CodeNode o) {
        return this.weight-o.weight;
    }

    @Override
    public String toString() {
        return "CodeNode{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }
    public void preOrder(){
        System.out.println(this);
        if (this.left!=null){
            this.left.preOrder();
        }
        if (this.right!=null){
            this.right.preOrder();
        }
    }


}


